# test.py
import torch
from torch.utils.data import DataLoader
from sar_cnn import SimpleCNN, MSTAR_Dataset  # 从 sar_cnn 导入所需内容
from torchvision import transforms
import os

# 指定计算设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据预处理
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=3),
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载测试数据集
test_dataset = MSTAR_Dataset(root_dir='./data/mstar/test', transform=transform)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 加载已保存的模型
num_classes = len(os.listdir('./data/mstar/train'))
model = SimpleCNN(num_classes=num_classes).to(device)
model.load_state_dict(torch.load('./output/sar_cnn_model.pth'))
model.eval()

# 测试模型
def test_model(model, test_loader):
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f"Test Accuracy: {100 * correct / total:.2f}%")

# 执行测试
test_model(model, test_loader)
